'''
Created on Nov 29, 2009

@author: mkiyer
'''

__version__ = '0.0.1'

import logging
import sys
import os
import operator

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import FormatStrFormatter
from matplotlib.patches import Polygon

from phdb import PileupDB
from veggie.db.sample.samplegroup import parse_samplegroups_xml
from veggie.io.bed import bed_reader


def plot_sample_coverage(chrom, start, end,
                         sample_groups, sample_coverages,
                         **kwargs):
    '''
    Plots individual sample coverage across the specified interval
    
    kwargs:
    axes: matplotlib Axes instance to plot into
    shade_start: start of region to shade
    shade_end: end of region to shade
    '''
    ax1 = kwargs.get('axes', plt.gca())     
    x = np.arange(0, end-start, 1)
    for sgroup in sample_groups:
        for s in sgroup.samples:
            ax1.plot(x, sample_coverages[s], color=sgroup.color, linewidth=1.0, label=s, alpha=0.75)
            ax1.annotate(s, xy=(x[0], np.max(sample_coverages[s])), xycoords='data',
                         size='xx-small', color=sgroup.color)

    shade_start = kwargs.get('shade_start', None)
    if shade_start is not None: 
        ax1.axvline(x=(shade_start - start), linewidth=2, color='black')
            
    ax1.grid(True)
    ax1.set_title("Individual samples", size='small')
    #ax1.set_title("%s:%d-%d" % (chrom, start, end), size='small')
    
def plot_boxplot_coverage(feature, sample_group_queries, 
                          **kwargs):
    '''
    sample_groups: list of SampleGroup objects.  the list order will be the
    order of the boxes    
    sample_coverages: dictionary keyed by sample name, values are 
    coverage data for that sample
    
    returns matplotlib Figure object
    
    kwargs:
    axes: matplotlib Axes instance to plot into
    title: title
    '''
    default_title = '%s:%d-%d' % (chrom, start, end)

    # extract from feature
    chrom, start, end = feature.chrom, feature.start, feature.en
    # lookup coverage
    boxplotdata = [query.coverage(chrom, start, end, rpkm=True)
                   for query in sample_group_queries]

    # setup plot
    ax1 = kwargs.get('axes', plt.gca())        

    # draw the boxes
    bp = ax1.boxplot(boxplotdata, notch=0, sym='+', vert=1, whis=1.5)
    plt.setp(bp['boxes'], color='black')
    plt.setp(bp['whiskers'], color='black')
    plt.setp(bp['fliers'], color='red', marker='+')

    # Add a horizontal grid to the plot, but make it very light in color
    # so we can use it for reading data values but not be distracting
    ax1.yaxis.grid(True, linestyle='-', which='major', color='lightgrey',
                  alpha=0.5)
    # Hide these grid behind plot objects
    ax1.set_axisbelow(True)
    
    # Set plot labels
    ax1.set_title('Sample groups', size='small')
    #ax1.set_title(kwargs.get('title', default_title))
    #ax1.set_xlabel('Sample group', size='x-small')
    ax1.set_ylabel('Coverage (RPKM)', size='small')

    # Now fill the boxes with desired colors
    boxColors = [squery.color for squery in sample_group_queries]
    numBoxes = len(boxplotdata)
    medians = range(numBoxes)
    for i in range(numBoxes):
        box = bp['boxes'][i]
        boxX = []
        boxY = []
        for j in range(5):
            boxX.append(box.get_xdata()[j])
            boxY.append(box.get_ydata()[j])
        boxCoords = zip(boxX, boxY)
        boxPolygon = Polygon(boxCoords, facecolor=boxColors[i])
        ax1.add_patch(boxPolygon)
        # Now draw the median lines back over what we just filled in
        med = bp['medians'][i]
        medianX = []
        medianY = []
        for j in range(2):
            medianX.append(med.get_xdata()[j])
            medianY.append(med.get_ydata()[j])
            ax1.plot(medianX, medianY, 'k')
            medians[i] = medianY[0]
        # Finally, overplot the sample averages, with horizontal alignment
        # in the center of each box
        ax1.plot([np.average(med.get_xdata())], [np.average(boxplotdata[i])],
                 color='w', marker='*', markeredgecolor='k')

    # Set the axes ranges and axes labels
    ax1.set_xlim(0.5, numBoxes+0.5)
    # TODO: setting ylim appears to screw things up
    bottom, top = ax1.get_ylim()
    #top = 1000
    #bottom = -5
    #ax1.set_ylim(bottom, top)
    labels = [squery.name for squery in sample_group_queries]
    xtickNames = plt.setp(ax1, xticklabels=labels)
    #plt.setp(xtickNames, rotation=45, fontsize=8)
    plt.setp(xtickNames, rotation=0, size='x-small')
    
    # Due to the Y-axis scale being different across samples, it can be
    # hard to compare differences in medians across the samples. Add upper
    # X-axis tick labels with the sample medians to aid in comparison
    # (just use two decimal places of precision)
    pos = np.arange(numBoxes)+1
    upperLabels = [str(np.round(s, 2)) for s in medians]
    weights = ['bold', 'semibold']
    for tick,label in zip(range(numBoxes),ax1.get_xticklabels()):
       k = tick % 2
       ax1.text(pos[tick], top-(top*0.05), upperLabels[tick],
            horizontalalignment='center', size='x-small', weight=weights[k],
            color=boxColors[tick])
    
    # Finally, add a basic legend
    ystart = 0.045
    yincrement = 0.02
    for squery in sample_group_queries:
        plt.figtext(0.80, ystart, squery.name,
                    backgroundcolor=squery.color, color='black', weight='roman',
                    size='x-small')
        ystart += yincrement
    plt.figtext(0.80, 0.015, '*', color='white', backgroundcolor='silver',
               weight='roman', size='medium')
    plt.figtext(0.815, 0.013, ' Average Value', color='black', weight='roman',
               size='x-small')


def plot_barplot(chrom, start, end,
                 sample_groups, sample_coverages,
                 **kwargs):
    '''
    sample_groups: list of SampleGroup objects.  the list order will be the
    order of the bar groups    
    sample_coverages: dictionary keyed by sample name, values are 
    coverage data for that sample
    
    returns matplotlib Figure object
    
    kwargs:
    axes: matplotlib Axes instance to plot into
    title: title
    '''
    default_title = '%s:%d-%d' % (chrom, start, end)
    # setup plot
    ax = kwargs.get('axes', plt.gca())    
    # construct boxplot data by making a list of vectors for
    # each sample group
    barlabels = []
    bardata = []
    barcolors = []
    for sgroup in sample_groups:
        barsamples = [(s, sample_coverages[s]) for s in sgroup.samples]
        barsamples = sorted(barsamples, key=operator.itemgetter(1), reverse=False)
        barlabels.extend([x[0] for x in barsamples])
        bardata.extend([x[1] for x in barsamples])
        barcolors.extend([sgroup.color for x in barsamples])

    ind = np.arange(len(bardata))
    width = 0.8
    ax.bar(ind, bardata, width, color=barcolors)
    # labels
    ax.set_ylabel('Coverage (RPKM)')
    ax.set_title("Individual samples", size='x-small')

    #ax.set_title('Scores by group and gender')
    ax.set_xticks(ind+width/2.0)
    ax.set_xticklabels(barlabels, rotation='vertical', size='xx-small')
    
    #ax.legend( (rects1[0], rects2[0]), ('Men', 'Women') )
    #autolabel(rects1)
    #autolabel(rects2)

#    def autolabel(rects):
#        # attach some text labels
#        for rect in rects:
#            height = rect.get_height()
#            ax.text(rect.get_x()+rect.get_width()/2., 1.05*height, '%d'%int(height),
#                    ha='center', va='bottom')


def plot_barplot2(chrom, start, end, sample_group_queries,
                 **kwargs):
    '''
    sample_groups: list of SampleGroup objects.  the list order will be the
    order of the bar groups    
    sample_coverages: dictionary keyed by sample name, values are 
    coverage data for that sample
    
    returns matplotlib Figure object
    
    kwargs:
    axes: matplotlib Axes instance to plot into
    title: title
    '''
    default_title = '%s:%d-%d' % (chrom, start, end)
    # setup plot
    ax = kwargs.get('axes', plt.gca())    

    # construct barplot data by making a list of vectors for
    # each sample group
    start_index = 0
    width = 0.8
    barticks = []
    for squery in sample_group_queries:
        querylabels = ['%s_%s' % (s, g) for s, g in zip(squery.samples, squery.groups)]
        arr = squery.coverage(chrom, start, end, rpkm=True)        
        querybars = [(s, x) for s, x in zip(querylabels, arr)]
        querybars = sorted(querybars, key=operator.itemgetter(1), reverse=False)
        # plot each group of bars with a different label so we can construct a legend
        barticks.extend([x[0] for x in querybars])        
        ind = np.arange(start_index, start_index + len(querybars))
        ax.bar(ind, [x[1] for x in querybars], width, color=squery.color, label=squery.name)
        start_index += len(querybars)
    # ticks
    ax.set_xticks(np.arange(len(barticks)) + width/2.0)
    ax.set_xticklabels(barticks, rotation='vertical', size='xx-small')
    # labels
    ax.set_ylabel('Expression (RPKM)')
    ax.set_title("Individual flowcells/lanes", size='x-small')

    ax.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
    #ax.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=2,
    #          mode="expand", borderaxespad=0.)
    #ax.legend(loc="best")
    return

    mylabels = ['%s_%s' % (s, g) for s, g in zip(squery.samples, squery.groups)]
    barsamples = [(s, x) for s, x in zip(mylabels, arr)]
    barsamples = sorted(barsamples, key=operator.itemgetter(1), reverse=False)
    barlabels.extend([x[0] for x in barsamples])
    bardata.extend([x[1] for x in barsamples])
    barcolors.extend([squery.color for x in barsamples])

        

    barlabels = []
    bardata = []
    barcolors = []
    for squery in sample_group_queries:
        arr = squery.coverage(chrom, start, end, rpkm=True)
        mylabels = ['%s_%s' % (s, g) for s, g in zip(squery.samples, squery.groups)]
        barsamples = [(s, x) for s, x in zip(mylabels, arr)]
        barsamples = sorted(barsamples, key=operator.itemgetter(1), reverse=False)
        barlabels.extend([x[0] for x in barsamples])
        bardata.extend([x[1] for x in barsamples])
        barcolors.extend([squery.color for x in barsamples])

    ind = np.arange(len(bardata))
    width = 0.8
    ax.bar(ind, bardata, width, color=barcolors)
    # labels
    ax.set_ylabel('Coverage (RPKM)')
    ax.set_title("Individual samples", size='x-small')
    #ax.set_title('Scores by group and gender')
    ax.set_xticks(ind+width/2.0)
    ax.set_xticklabels(barlabels, rotation='vertical', size='xx-small')

    ax.legend( (rects1[0], rects2[0]), ('Men', 'Women') )


    # Finally, add a basic legend
    ystart = 0.045
    yincrement = 0.02
    for squery in sample_group_queries:
        plt.figtext(0.80, ystart, squery.name,
                    backgroundcolor=squery.color, color='black', weight='roman',
                    size='x-small')
        ystart += yincrement
    plt.figtext(0.80, 0.015, '*', color='white', backgroundcolor='silver',
               weight='roman', size='medium')
    plt.figtext(0.815, 0.013, ' Average Value', color='black', weight='roman',
               size='x-small')    


def plot_coverage(feature, sample_group_queries, savefile,
                  **kwargs):
    '''
    kwargs:
    figsize - size of figure (tuple)
    title - title of plot (string)
    '''
    # TODO: convert to a GFF3-like feature format instead of BED
    chrom, start, end = feature.chrom, feature.start, feature.end

    # TODO: move this to a global location for plotting settings
    major_formatter = FormatStrFormatter('%d')

    # setup figure
    default_suptitle = '%s:%d-%d' % (chrom, start, end)
    default_figsize = (10,8)
    fig = plt.figure(figsize=kwargs.get('figsize', default_figsize))
    fig.suptitle(kwargs.get('title', default_suptitle))

    # show coverage
    ax1 = fig.add_subplot(311)
    x = np.arange(start, end, 1)
    for query in sample_group_queries:
        # fetch the coverage from the database
        arr = query.coverage(chrom, start, end, norm=True)
        # plot each sample
        for i in xrange(arr.shape[0]):
            label = '%s: %s' % (query.samples[i], query.groups[i])
            ax1.plot(x, arr[i], color=query.color, linewidth=1.0, label=label, alpha=0.75)
            ax1.annotate(label, xy=(x[0], np.max(arr[i])), xycoords='data',
                         size='xx-small', color=query.color)
    ax1.xaxis.set_major_formatter(major_formatter)
    ax1.set_xlim(start, end)

    # show splice junctions
    ax2 = fig.add_subplot(312)
    ypos = 0
    ybarw = 10
    yticks = []
    yticklabels = []
    for query in sample_group_queries:
        # fetch the coverage from the database
        query_juncs = query.junctions(chrom, start, end, norm=True)
        for i, junc_list in enumerate(query_juncs):
            ax2.broken_barh([(j[0], j[1]-j[0]) for j in junc_list],
                            (ypos, ybarw - 1),
                            facecolors=query.color)
            yticks.append(ypos + int(ybarw/2.0))
            label = '%s: %s' % (query.samples[i], query.groups[i])
            yticklabels.append(label)
            ypos += ybarw
    ax2.set_yticks(yticks)
    ax2.set_yticklabels(yticklabels, size='xx-small')
    ax2.xaxis.set_major_formatter(major_formatter)
    ax2.set_xlim(start, end)
    
    # show the feature model
    # TODO: query the feature database
    ax3 = fig.add_subplot(313)
    if feature.blockCount != None:
        for i in xrange(feature.blockCount):
            ax3.axvspan(start + feature.blockStarts[i],
                        start + feature.blockStarts[i] + feature.blockSizes[i],
                        facecolor='g', alpha=0.5)
    ax3.xaxis.set_major_formatter(major_formatter)
    ax3.set_xlim(start, end)    
    #plt.show()
    plt.savefig(savefile)
    plt.close(fig)


def plot_feature(feature, sample_group_queries, savefile, **kwargs):
    '''
    kwargs:
    figsize - size of figure (tuple)
    title - title of plot (string)
    shade_start - instructs certain plots to shade a region
    shade_end - instructs certain plots to shade a region
    '''
    # extract info from feature
    chrom, start, end = feature.chrom, feature.start, feature.end

    # setup figure
    default_suptitle = 'Expression for %s:%d-%d' % (chrom, start, end)
    default_figsize = (15,10)
    #default_window_title = 'Boxplot'
    fig = plt.figure(figsize=kwargs.get('figsize', default_figsize))
    fig.suptitle(kwargs.get('title', default_suptitle))
    #fig.canvas.set_window_title(kwargs.get('window_title', default_window_title))
    fig.subplots_adjust(left=0.05, right=0.75, top=0.95, bottom=0.2,
                        wspace=0.1, hspace=0.25)

    ax1 = fig.add_subplot(111)
    plot_barplot2(chrom, start, end, sample_group_queries, axes=ax1)
    plt.savefig(savefile)
    plt.close(fig)
    return
    
    # get coverage data
    sample_coverages = covdb.calc_avg_cov_arrays(chrom, start, end, sample_groups)
    summed_coverages = {}
    for s, covarray in sample_coverages.iteritems():
        summed_coverages[s] = covdb.calc_rpkm_array(covarray)



    ax1 = fig.add_subplot(221)
    ax2 = fig.add_subplot(223)
    ax3 = fig.add_subplot(222)
    ax4 = fig.add_subplot(224)
    # Top plot shows individual samples
    oncoplot.plot_sample_coverage(chrom, start, end, 
                                  sample_groups, sample_coverages,
                                  axes=ax1,
                                  **kwargs)
    # Bottom plot shows average for groups
    oncoplot.plot_sample_group_average_coverage(chrom, start, end, 
                                                sample_groups, 
                                                sample_coverages,
                                                axes=ax2,
                                                **kwargs)
    # Barplot
    oncoplot.plot_barplot(chrom, start, end,
                          sample_groups, summed_coverages,
                          axes=ax3)
    # Boxplot
    oncoplot.plot_boxplot(chrom, start, end, 
                          sample_groups, summed_coverages,
                          axes=ax4)
    plt.savefig('%s_%d-%d_expr.png' % (chrom, start, end))
    #plt.show()
    plt.close()
    return


if __name__ == '__main__':
    # setup logging
    logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s',
                        level=logging.DEBUG)
    logger = logging.getLogger(__name__)
    logger.info("vizdb version %s" % (__version__))
    logger.info("-----------------------")

    # parse command line arguments
    samples_xmlfile = sys.argv[1]
    interval_file = sys.argv[2]

    # connect to pileup database
    db_path = '/home/mkiyer/exper/prostate_rnaseq_2010_Jan_13.phd'
    phdb = PileupDB(db_path, 'r')

    # connect to gene database
    
    # load samples and construct database queries
    samples = []
    sample_group_queries = []
    for sgroup in parse_samplegroups_xml(samples_xmlfile):
        samples.extend(sgroup.samples)
        # construct queries and attach sample group attributes
        myquery = phdb.query(sgroup.samples, qc_filter=True)
        myquery.name = sgroup.name
        myquery.color = sgroup.color
        sample_group_queries.append(myquery)

    # load intervals
    for feature in bed_reader(open(interval_file), cols=4):
        savefile = '%s_%s_%d-%d.png' % (feature.name, feature.chrom, feature.start, feature.end)
        # plot coverage for each interval
        #plot_coverage(feature, sample_group_queries, savefile)
        plot_feature(feature, sample_group_queries, savefile)

    phdb.close()
